package com.smart.milvus.conditions.search;

import lombok.Getter;
import lombok.Setter;
import lombok.experimental.Accessors;
import com.smart.milvus.bean.MilvusModel;
import com.smart.milvus.conditions.AbstractWrapper;
import com.smart.milvus.conditions.segments.MergeSegments;
import com.smart.milvus.model.index.IIndexExtra;
import com.smart.milvus.base.func.SFunction;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

@Setter
@Getter
@Accessors(chain = true)
public class SearchWrapper<T extends MilvusModel<? extends MilvusModel<?>>>
        extends AbstractWrapper<T, String, SearchWrapper<T>>
        implements BaseSearchWrapper<T, String, SearchWrapper<T>> {

    /**
     * 额外的索引查询参数
     * Search parameter(s) specific to the index.
     * See <a href="https://milvus.io/docs/index.md">Vector Index</a> for more information.
     */
    private IIndexExtra extra;
    private String indexName;
    private String indexColumn;
    /**
     * 查询最相似的多少条数据
     * Number of the most similar results to return.
     */
    private Integer topK = 10;
    /**
     * 指定要检索的向量列
     */
    private String vectorColumn;
    /**
     * 指定输入向量
     */
    private List<Float> vectorValue;

    public SearchWrapper() {
        this((T) null);
    }

    public SearchWrapper(T entity) {
        super.setEntity(entity);
        super.initNeed();
    }

    public SearchWrapper(Class<T> entityClass) {
        super.setEntityClass(entityClass);
        super.initNeed();
    }

    /**
     * 非对外公开的构造方法,只用于生产嵌套 Expr
     *
     * @param entityClass 本不应该需要的
     */
    private SearchWrapper(T entity, Class<T> entityClass, AtomicInteger paramNameSeq,
                          MergeSegments mergeSegments) {
        super.setEntity(entity);
        super.setEntityClass(entityClass);
        this.paramNameSeq = paramNameSeq;
        this.expression = mergeSegments;
    }

    private SearchWrapper(T entity, Class<T> entityClass, AtomicInteger paramNameSeq,
                          MergeSegments mergeSegments, IIndexExtra extra, Integer topK,
                          String vectorColumn, List<Float> vectorValue) {
        super.setEntity(entity);
        super.setEntityClass(entityClass);
        this.paramNameSeq = paramNameSeq;
        this.expression = mergeSegments;
        this.extra = extra;
        this.topK = topK;
        this.vectorColumn = vectorColumn;
        this.vectorValue = vectorValue;
    }


    @Override
    public SearchWrapper<T> vector(String column, List<Float> value) {
        this.vectorColumn = column;
        this.vectorValue = value;
        return this;
    }

    @Override
    public SearchWrapper<T> index(String indexName, String indexColumn, IIndexExtra indexExtra) {
        this.indexName = indexName;
        this.indexColumn = indexColumn;
        this.extra = indexExtra;
        return this;
    }

    /**
     * 用于生成嵌套 Expr
     * <p>
     * 故 ExprSelect 不向下传递
     * </p>
     */
    @Override
    protected SearchWrapper<T> instance() {
        return new SearchWrapper<>(getEntity(), getEntityClass(), paramNameSeq, new MergeSegments());
    }

    @Override
    public void clear() {
        super.clear();
    }
}
